# -*- coding: utf-8 -*-
"""
#-------------------------------------------------------------------------------
# Nom:         traitementModelePQ
#
# Auteur:      Alain Gauthier
#
# Créé le:     16/08/2013
# Copyright:   (c) eaucéa 2014
# Licence:     GNU Public Licence (GPL)

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/

Calage et lancement du modèle GR2M à échelle régionale:
    - Calage à partir d'une liste de stations de mesure listées dans un fichier CSV
    - Régionalisation du modèle selon les possibilités proposées
    - Lancement du modèle sur des points sélectionnés représentants des exutoires de Bassins Versant de calcul

Ce script doit être lancé depuis l'interface du plugin 'ModelePluieDebit' sous QGIS

#-------------------------------------------------------------------------------
"""

# modules système
import os
import sys
import traceback
import datetime
# modules externes
import numpy as np
import numpy.linalg as LA
import scipy.stats as stat
# modules QGIS
from qgis.core import *
# modules GUI
from PyQt4 import QtCore, QtGui
# modules internes
import constantesGeo as cg
import importCalage
import analyseChroniques
import optimGR
import gr2m
import typeCritereNash
import affChroniques

# constante de nom de fichier log
__fichier_log__= os.path.join(os.getenv('TEMP'),"modelepluiedebit_ModelePQ.log")

class TraitementModelePQ(object):
    """
    Classe servant à contrôler les paramètres du modèle, à lancer le calage et les simulations.
    """
    def __init__(self,
                 AcheminChroniquesPluie,
                 AcheminChroniquesEvapo,
                 AdtPluie,
                 AdtEvapo):
        """
        Constructeur:
            - AcheminChroniquesPluie  les chroniques de pluie
            - AcheminChroniquesEvapo  les chroniques d'Evapotranspiration
            - AdtPluie                Pas de temps des chroniques de pluie
            - AdtEvapo                Pas de temps des chroniques d'Evapotranspiration
        """
        self.mettreAJourDonnees(AcheminChroniquesPluie,AcheminChroniquesEvapo,AdtPluie,AdtEvapo)
        self.setprogressBar(None)
        # initialisation des données
        self.chroniquesPluie = dict()
        self.chroniquesEvapo = dict()
        self.chroniquesQmesure = dict()
        self.resultatOptim = dict()
        self.critereOptim = dict()
        self.dicoStations=dict()

    #--------------------------------------------------------------------------#

    def mettreAJourDonnees(self,
                           AcheminChroniquesPluie,
                           AcheminChroniquesEvapo,
                           AdtPluie,
                           AdtEvapo):
        """
        Mise à jour de données utilisateur:
            - AcheminChroniquesPluie  les chroniques de pluie
            - AcheminChroniquesEvapo  les chroniques d'Evapotranspiration
            - AdtPluie                Pas de temps des chroniques de pluie
            - AdtEvapo                Pas de temps des chroniques d'Evapotranspiration
        """
        self.cheminChroniquesPluie = AcheminChroniquesPluie
        self.cheminChroniquesEvapo = AcheminChroniquesEvapo
        self.dtPluie = AdtPluie
        self.dtEvapo = AdtEvapo

    #--------------------------------------------------------------------------#

    def setprogressBar(self,AprogressBar):
        """
        gestion de la barre de progression
        """
        self.progressBar = AprogressBar

    #--------------------------------------------------------------------------#

    def ajusterProgressBar(self,Avaleur,Amax):
        """
        ajustement de la jauge en utilisant le rapport Avaleur/Amax
        """
        if Amax==0:
            self.progressBar.setMaximum(0)
            self.progressBar.setValue(0)
        else:
            proportion=100*Avaleur/Amax
            self.progressBar.setValue(proportion)
            QtGui.QApplication.processEvents()

    #--------------------------------------------------------------------------#

    def chargerStationsCalage(self, aNomFichier):
        """
        Lit le fichier des stations servant au calage et initialise les données de calage venant du fichier aNomFichier
        Renvoie une dictionnaire contenant les données des stations à caler
        """
        resultat = importCalage.chargerStations(aNomFichier)
        return resultat

    #--------------------------------------------------------------------------#

    def dtChronique(self,Achronique):
        """
        Calcule et renvoie le dt associé à la chronique Achronique
        si la chronique est mensuelle alors dt=31 jours.
        """
        tab_dt = Achronique['temps']
        date0 = tab_dt[0]
        date1 = tab_dt[1]
        dt = date1-date0
        # cas particulier du pas de temps mensuel : fixer dt=31 jours
        if date1.month != date0.month:
            dt = datetime.timedelta(31)
        return dt

    #--------------------------------------------------------------------------#

    def lancerCalage(self,
                     AfichierStationCalage,
                     AfichierQmesure,
                     AcritereEfficacite,
                     AparamOptim):
        """
        Lancement du calage avec les paramètres fixés :
            - AfichierStationCalage   Fichier de calage contenant chaque station sélectionnée
            - AfichierQmesure         Fichier de Q mesuré
            - AcritereEfficacite      Critère d'efficacité choisi pour le calage
            - AparamOptim             Paramètres à optimiser pour le calage, les autres sont fixés
        """
        self.fichierLog = open(__fichier_log__,'a')
        try:
            try:
                # initialisation des données
                self.chroniquesPluie = dict()
                self.chroniquesEvapo = dict()
                self.chroniquesQmesure = dict()
                self.resultatOptim = dict()
                self.critereOptim = dict()

                # chargement du fichier de calage
                self.fichierLog.write('chargement du fichier de calage : '+AfichierStationCalage+'\r\n')
                self.dicoStations = self.chargerStationsCalage(AfichierStationCalage)

                nbStations = len(self.dicoStations.keys())
                cpt=0
                # Boucle sur chaque station pour charger les chroniques associées
                for idStation in self.dicoStations:
                    # données de la station courante
                    donneesStation = self.dicoStations[idStation]
                    self.fichierLog.write('calage de la station de : '+repr(donneesStation['nom_station'])+'\r\n')

                    # identifiants des bassins versant des stations
                    bvStation = donneesStation['id_BV']
                    # chargement des chroniques
                    self.chroniquesPluie[idStation] = analyseChroniques.construire(bvStation,
                                                                                   self.cheminChroniquesPluie,
                                                                                   self.dtPluie)
                    self.chroniquesEvapo[idStation] = analyseChroniques.construire(bvStation,
                                                                                   self.cheminChroniquesEvapo,
                                                                                   self.dtEvapo)
                    self.chroniquesQmesure[idStation] = analyseChroniques.chargerChroniqueQ(idStation,
                                                                                            AfichierQmesure)

                    # traitement spécifique des données GLDAS de la chronique Evapotranspiration :
                    # les valeurs sont des moyennes sur des durées de 3h et
                    # la chronique est construite sur les dates au pas de temps de la pluie
                    self.chroniquesEvapo[idStation] = analyseChroniques.traiterDonneesGLDAS(self.chroniquesEvapo[idStation],
                                                                                            self.dtEvapo,
                                                                                            self.chroniquesPluie[idStation]['temps'])

                    # Préparation des chroniques pour les simulations :
                    # -------------------------------------------------
                    # la chronique mesurée doit être au même pas de temps et avec la même unité que la pluvio et l'évapo
                    # les chroniquest sont alignées afin de travailler sur la même ref d'indices relatifs

                    # chronique de débit mesuré (en mm)
                    tab_dtPluie = self.chroniquesPluie[idStation]['temps']
                    dtPluie = self.dtChronique(self.chroniquesPluie[idStation])

                    LQmes = analyseChroniques.debitEn_mm(self.chroniquesQmesure[idStation],
                                                         dtPluie,
                                                         donneesStation['aire'])

                    # alignement du début des chroniques :
                    self.chroniquesEvapo[idStation], \
                    self.chroniquesPluie[idStation], \
                    self.chroniquesQmesure[idStation] = analyseChroniques.alignerChroniques(self.chroniquesEvapo[idStation],
                                                                                            self.chroniquesPluie[idStation],
                                                                                            LQmes)

                    # lancement du calage spécifique à la station
                    self.resultatOptim[idStation] = optimGR.lancerOptimisation(AcritereEfficacite,
                                                                               self.chroniquesPluie[idStation],
                                                                               self.chroniquesEvapo[idStation],
                                                                               self.chroniquesQmesure[idStation],
                                                                               donneesStation,
                                                                               AparamOptim,
                                                                               self.fichierLog)

                    # lancement de la simulation et calcul des critères de Nash et Sutcliffe
                    xParam=np.array(self.resultatOptim[idStation])
                    lPluie=np.array(self.chroniquesPluie[idStation]['valeur'])
                    levapo=np.array(self.chroniquesEvapo[idStation]['valeur'])
                    lniveauIni= np.array([0.5,0.5])
                    # simulation GR2M
                    if dtPluie.days == 31:
                        resultatGR=gr2m.gr2m(xParam,lPluie,levapo,lniveauIni)
                    else:
                        raise NotImplementedError("GR4J n'est pas implémenté.")
                    # critères de Nash et Sutcliffe sur l'intervalle d'évaluation
                    dt_station = [analyseChroniques.decoder_date(tps) for tps in donneesStation['duree_eval']]
                    Ldt_eval = np.array([tab_dtPluie.index(dt_station[0]),tab_dtPluie.index(dt_station[1])])
                    LQ_GR = np.array(resultatGR['Q'][Ldt_eval[0]:Ldt_eval[1]])
                    LQ_ref= np.array(self.chroniquesQmesure[idStation]['valeur'][Ldt_eval[0]:Ldt_eval[1]])
                    self.critereOptim[idStation]=[]
                    for critere in range(len(typeCritereNash.listeCritereNash)):
                        Nash = optimGR.calculerNash(LQ_GR,LQ_ref,critere)
                        self.critereOptim[idStation].append(Nash)

                    # mise à jour de progressBar
                    cpt+=1
                    self.ajusterProgressBar(cpt,nbStations)
            except Exception, e:
                exc_type, exc_value, exc_traceback = sys.exc_info()
                traceback.print_exc(None,sys.stdout)
                self.fichierLog.write('-'*80+'\n')
                self.fichierLog.write('exception : ' + str(e)+'\n')
                self.fichierLog.write(repr(traceback.extract_tb(exc_traceback)))
                self.fichierLog.write('\n'+'-'*80+'\n')
        finally:
            self.fichierLog.close()

    #--------------------------------------------------------------------------#

    def lancerSimulation(self,AIDcoucheBVcalcul):
        """
        lancement de la simulation du modèle à échelle régionale pour les BV sélectionnés de la couche AIDcoucheBVcalcul
        Renvoie un dictionnaire contenant pour chaque BV de calcul :
            - la chronique de pluie
            - la chronique d'évapotranspiration
            - la chronique de débit simulé (en mm)
            - l'aire du BV
            - la valeur des paramètres calculés par le modèle régional pour la simulation
        """

        # résultat
        resultat=dict()
        self.fichierLog = open(__fichier_log__,'a')
        try:
            try:
                self.fichierLog.write('début simulation\r\n')
                # BV sélectionnés
                coucheBV=QgsMapLayerRegistry.instance().mapLayer(str(AIDcoucheBVcalcul))
                ensemble_selec = coucheBV.selectedFeatures()

                # resolution des méthodes de régionalisation (calcul des courbes de régression)
                # à partir des paramètres obtenus par calage:
                ParamRegional=self.resoudreRegionalisation()

                # boucle sur les points sélectionnés
                cpt=0
                nbSelec=len(ensemble_selec)
                for entiteBV in ensemble_selec:
                    # attributs du BV
                    attributs = entiteBV.attributes()
                    id_bv = str(attributs[coucheBV.fieldNameIndex(cg.__ID_ND__)])
                    penteBV = attributs[coucheBV.fieldNameIndex(cg.__ID_SLP__)]
                    aire = attributs[coucheBV.fieldNameIndex(cg.__ID_AIRE__)]

##                self.fichierLog.write('id_bv = ' + str(id_bv)+'\r\n')
##                self.fichierLog.write('penteBV = ' + str(penteBV)+'\r\n')
##                self.fichierLog.write('aire = ' + str(aire)+'\r\n')

                    # initialisation
                    resultat[id_bv]=dict()
                    resultat[id_bv]['aire']=aire

                    # chroniques de pluie et évapo sur le bv de calcul
                    chroniquesPluie = analyseChroniques.construire(id_bv,
                                                                   self.cheminChroniquesPluie,
                                                                   self.dtPluie)
                    chroniquesEvapo = analyseChroniques.construire(id_bv,
                                                                   self.cheminChroniquesEvapo,
                                                                   self.dtEvapo)
                    # traitement spécifique des données GLDAS de la chronique Evapotranspiration :
                    # les valeurs sont des moyennes sur des durées de 3h et
                    # la chronique est construite sur les dates au pas de temps de la pluie
                    chroniquesEvapo = analyseChroniques.traiterDonneesGLDAS(chroniquesEvapo,
                                                                            self.dtEvapo,
                                                                            chroniquesPluie['temps'])

                    # en sortie
                    resultat[id_bv]['pluie']=chroniquesPluie
                    resultat[id_bv]['evapo']=chroniquesEvapo

                    # pour la simulation
                    lPluie=np.array(chroniquesPluie['valeur'])
                    levapo=np.array(chroniquesEvapo['valeur'])
                    lniveauIni= np.array([0.5,0.5])

                    # simulation GR2M ou GR4J selon le pas de temps
                    # simulation GR2M/GR4J
                    dtP=self.dtChronique(chroniquesPluie)

                    # régionalisation : calcul des paramètres du modèle pour les points sélectionnés
                    xParam=self.regionaliserGR(dtP,penteBV,ParamRegional)

                    self.fichierLog.write('BV n° : ' + str(cpt)+'\r\n')
                    self.fichierLog.write('xParam = ' + str(xParam)+'\r\n')

                    resultat[id_bv]['xParam']=xParam

                    if dtP.days == 31:
                        resultatGR=gr2m.gr2m(xParam,lPluie,levapo,lniveauIni)
                    else:
                        raise NotImplementedError("GR4J n'est pas implémenté.")
                    # en sortie
                    resultat[id_bv]['Qsimu']= resultatGR['Q']

                    # mise à jour de progressBar
                    cpt+=1
                    self.ajusterProgressBar(cpt,nbSelec)
                # fin de boucle
            except Exception, e:
                exc_type, exc_value, exc_traceback = sys.exc_info()
                traceback.print_exc(None,sys.stdout)
                self.fichierLog.write('-'*80+'\n')
                self.fichierLog.write('exception : ' + str(e)+'\n')
                self.fichierLog.write(repr(traceback.extract_tb(exc_traceback)))
                self.fichierLog.write('\n'+'-'*80+'\n')
        finally:
            self.fichierLog.close()
            # renvoie le résultat
            return resultat

    #--------------------------------------------------------------------------#

    def calculStat(self,AresultatSimu):
        """
        Calcul du QMNA5 et du module sur les chroniques
        Ajustement avec une loi de Galton (log-normale)
        """
        # résultat
        resultat = dict()
        self.fichierLog = open(__fichier_log__,'a')
        try:
            try:
                cpt=0
                nbBv=len(AresultatSimu.keys())
                # boucle sur les chroniques simulées
                for id_bv in AresultatSimu:
                    self.fichierLog.write('id_bv = ' + str(id_bv)+'\r\n')
                    # initialisation resultat
                    resultat[id_bv]=dict()
                    # chronique en en m3/s
                    Q=np.array(analyseChroniques.debit_enM3S(AresultatSimu[id_bv]['Qsimu'],
                                                             AresultatSimu[id_bv]['pluie']['temps'],
                                                             AresultatSimu[id_bv]['aire']))
                    # module (sur un multiple de 12 mois)
                    ideb=0
                    ifin=12
                    moyenneAnnuelle=[]
                    QMNA=[]
                    while ifin<=Q.size:
                        moyenneAnnuelle.append(np.mean(Q[ideb:ifin]))
                        QMNA.append(np.min(Q[ideb:ifin]))
                        ideb=ifin
                        ifin+=12
                    moyenneAnnuelle=np.array(moyenneAnnuelle)
                    QMNA=np.array(QMNA)
                    resultat[id_bv]['Module']=np.mean(moyenneAnnuelle)
                    # ln(QMNA)
                    lnQMNA=np.log(QMNA)
                    # loi de Galton
                    moy=np.mean(lnQMNA)
                    std=np.std(lnQMNA)
                    # QMNA5 (donnée quinquénale sèche)
                    lnQMNA5=stat.norm.ppf(np.array([0.2]),loc=moy,scale=std)
                    #resultat
                    resultat[id_bv]['QMNA5']=np.exp(lnQMNA5[0])

                    self.fichierLog.write('QMNA5 = ' + str(resultat[id_bv]['QMNA5'])+'\r\n')
                    self.fichierLog.write('Module = ' + str(resultat[id_bv]['Module'])+'\r\n')
                    # mise à jour de progressBar
                    cpt+=1
                    self.ajusterProgressBar(cpt,nbBv)
                # fin
            except Exception, e:
                exc_type, exc_value, exc_traceback = sys.exc_info()
                traceback.print_exc(None,sys.stdout)
                self.fichierLog.write('-'*80+'\n')
                self.fichierLog.write('exception : ' + str(e)+'\n')
                self.fichierLog.write(repr(traceback.extract_tb(exc_traceback)))
                self.fichierLog.write('\n'+'-'*80+'\n')
        finally:
            self.fichierLog.close()
            return resultat

    #--------------------------------------------------------------------------#

    def resoudreRegionalisation(self):
        """
        Résolution des courbes de régression pour l'estimation de paramètres régionaux
        Renvoie pour chaque paramètres les coefficients nécessaire à son calcul à échelle régionale
        resultat['X1']=(a,b) pour calculer X1=a*pente+b
        resultat['X2']=moyenne(X2optimal)
        resultat['X3']=60 pour GR2M
        resultat['X4']=None pour GR2M
        """
        # Pour X1 dans le cas GR2M, on cherche [a,b] satisfaisant A*[a,b]= X1calage au sens des moindres carrés
        # où :
        #   X1calage est donné par les X1 calés pour les stations de mesure
        #   A = [pente 1] est la matrice calculée avec les pentes des bv
        resultatGR2M=dict()
        X1Optim=[]
        penteOptim=[]
        for idStation in self.resultatOptim:
            X1Optim.append(self.resultatOptim[idStation][0])
            penteOptim.append(float(self.dicoStations[idStation]['pente']))

        A=np.vstack([np.array(penteOptim), np.ones(len(penteOptim))]).T

        X1calage=np.array(X1Optim)
        # résolution au sens des moindres carrés, les valeurs de 1 seront calculées sous la forme
        # X1=a*pente+b
        a,b= LA.lstsq(A, X1calage)[0]
        resultatGR2M['X1'] = (a,b)

        # pour X2 dans le cas GR2M, on prend la moyenne des valeurs optimisées en expluant les extrêmes
        X2optim=[]
        for idStation in self.resultatOptim:
            X2optim.append(self.resultatOptim[idStation][1])
        X2optim=np.array(X2optim)
        X2min=np.min(X2optim)
        X2max=np.max(X2optim)
        resultatGR2M['X2'] = np.mean(X2optim[(X2optim>X2min)&(X2optim<X2max)])

        # pour X3 dans le cas GR2M, X3=60 mm
        resultatGR2M['X3'] = 60

        # pour X4 dans le cas GR2M pas de valeur
        resultatGR2M['X4'] = None

        # fin
        return resultatGR2M

    #--------------------------------------------------------------------------#

    def regionaliserGR(self,Adt,ApenteBV,AParamRegional):
        """
        Régionalisation des paramètres du modèle GR en appliquant les coef calculés dans AParamRegional
        renvoie les paramètres X du modèle GR2M calculés pour le BV de pente ApenteBV
        """
        # initialisation
        xRegional=np.ones(4)
        # test
        if len(self.resultatOptim.keys())==0:
            QtGui.QMessageBox.information(None,QtGui.QApplication.translate('widgetDonnees','Information'), \
                QtGui.QApplication.translate('widgetDonnees',u'Lancer le calage du modèle avant une simulation'))

        # cas du modèle GR2M
        if Adt.days == 31:
            # X1 obtenu par régression linéaire/penteBV en excluant les valeurs extrêmes
            xRegional[0]=AParamRegional['X1'][0]*ApenteBV+AParamRegional['X1'][1]
            # X2 obtenu par la moyenne des valeurs optimales en excluant les valeurs extrêmes
            xRegional[1]=AParamRegional['X2']
            # X3 fixé = 60 mm
            xRegional[2]=AParamRegional['X3']
        # cas du modèle GR4J
        else:
            raise NotImplementedError("GR4J n'est pas implémenté.")
        # résultat
        return xRegional

    #--------------------------------------------------------------------------#
    #--------------------------------------------------------------------------#

if __name__ == '__main__':
    raise EnvironmentError("module devant être exécuté sous console/extension python QGIS")
